'''
This demo shows the communication interface of the MR813 motion control board based on LCM.
- robot_control_cmd_lcmt.py
- file_send_lcmt.py
- Gait_Def_moonwalk.toml
- Gait_Params_moonwalk.toml
- Usergait_List.toml
'''
import lcm
import sys
import time
import toml
import copy
import math
from robot_control_cmd_lcmt import robot_control_cmd_lcmt
from file_send_lcmt import file_send_lcmt

# 定义机器人控制命令的默认参数
robot_cmd = {
    'mode': 0, 'gait_id': 0, 'contact': 0, 'life_count': 0,
    'vel_des': [0.0, 0.0, 0.0],
    'rpy_des': [0.0, 0.0, 0.0],
    'pos_des': [0.0, 0.0, 0.0],
    'acc_des': [0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
    'ctrl_point': [0.0, 0.0, 0.0],
    'foot_pose': [0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
    'step_height': [0.0, 0.0],
    'value': 0, 'duration': 0
}

def main():
    # 初始化 LCM 客户端，用于发送控制命令和用户自定义步态文件
    lcm_cmd = lcm.LCM("udpm://239.255.76.67:7671?ttl=255")
    lcm_usergait = lcm.LCM("udpm://239.255.76.67:7671?ttl=255")
    usergait_msg = file_send_lcmt()  # 用户自定义步态文件消息
    cmd_msg = robot_control_cmd_lcmt()  # 机器人控制命令消息

    try:
        # 加载步态参数配置文件
        steps = toml.load("Gait_Params_moonwalk.toml")
        full_steps = {'step': [robot_cmd]}  # 初始化完整步态参数
        k = 0
        for i in steps['step']:
            cmd = copy.deepcopy(robot_cmd)  # 深拷贝默认参数
            cmd['duration'] = i['duration']  # 设置持续时间
            if i['type'] == 'usergait':
                # 设置用户自定义步态参数
                cmd['mode'] = 11  # LOCOMOTION 模式
                cmd['gait_id'] = 110  # USERGAIT 步态 ID
                cmd['vel_des'] = i['body_vel_des']  # 身体速度
                cmd['rpy_des'] = i['body_pos_des'][0:3]  # 身体姿态
                cmd['pos_des'] = i['body_pos_des'][3:6]  # 身体位置
                cmd['foot_pose'][0:2] = i['landing_pos_des'][0:2]  # 脚部姿态
                cmd['foot_pose'][2:4] = i['landing_pos_des'][3:5]
                cmd['foot_pose'][4:6] = i['landing_pos_des'][6:8]
                cmd['ctrl_point'][0:2] = i['landing_pos_des'][9:11]  # 控制点
                # 计算步高
                cmd['step_height'][0] = math.ceil(i['step_height'][0] * 1e3) + math.ceil(i['step_height'][1] * 1e3) * 1e3
                cmd['step_height'][1] = math.ceil(i['step_height'][2] * 1e3) + math.ceil(i['step_height'][3] * 1e3) * 1e3
                cmd['acc_des'] = i['weight']  # 期望加速度
                cmd['value'] = i['use_mpc_traj']  # 使用 MPC 轨迹
                cmd['contact'] = math.floor(i['landing_gain'] * 1e1)  # 接触增益
                cmd['ctrl_point'][2] = i['mu']  # 摩擦系数
            if k == 0:
                full_steps['step'] = [cmd]  # 初始化完整步态参数
            else:
                full_steps['step'].append(cmd)  # 添加步态参数
            k += 1

        # 将完整步态参数保存到新文件
        f = open("Gait_Params_moonwalk_full.toml", 'w')
        f.write("# Gait Params\n")
        f.writelines(toml.dumps(full_steps))
        f.close()

        # 发送用户自定义步态文件
        file_obj_gait_def = open("Gait_Def_moonwalk.toml", 'r')
        file_obj_gait_params = open("Gait_Params_moonwalk_full.toml", 'r')
        usergait_msg.data = file_obj_gait_def.read()  # 读取步态定义文件
        lcm_usergait.publish("user_gait_file", usergait_msg.encode())  # 发送步态定义文件
        time.sleep(0.5)
        usergait_msg.data = file_obj_gait_params.read()  # 读取步态参数文件
        lcm_usergait.publish("user_gait_file", usergait_msg.encode())  # 发送步态参数文件
        time.sleep(0.1)
        file_obj_gait_def.close()
        file_obj_gait_params.close()

        # 加载用户自定义步态列表
        user_gait_list = open("Usergait_List.toml", 'r')
        steps = toml.load(user_gait_list)
        for step in steps['step']:
            # 设置控制命令参数
            cmd_msg.mode = step['mode']
            cmd_msg.value = step['value']
            cmd_msg.contact = step['contact']
            cmd_msg.gait_id = step['gait_id']
            cmd_msg.duration = step['duration']
            cmd_msg.life_count += 1  # 更新生命计数，确保命令生效
            for i in range(3):
                cmd_msg.vel_des[i] = step['vel_des'][i]  # 期望速度
                cmd_msg.rpy_des[i] = step['rpy_des'][i]  # 期望姿态
                cmd_msg.pos_des[i] = step['pos_des'][i]  # 期望位置
                cmd_msg.acc_des[i] = step['acc_des'][i]
                cmd_msg.acc_des[i+3] = step['acc_des'][i+3]  # 期望加速度
                cmd_msg.foot_pose[i] = step['foot_pose'][i]  # 脚部姿态
                cmd_msg.ctrl_point[i] = step['ctrl_point'][i]  # 控制点
            for i in range(2):
                cmd_msg.step_height[i] = step['step_height'][i]  # 步高
            # 发送控制命令
            lcm_cmd.publish("robot_control_cmd", cmd_msg.encode())
            print('robot_control_cmd lcm publish mode :', cmd_msg.mode, "gait_id :", cmd_msg.gait_id, "msg.duration=", cmd_msg.duration)
            time.sleep(0.1)  # 短暂休眠

        # 发送心跳信号，保持通信
        for i in range(75):  # 15 秒心跳信号
            lcm_cmd.publish("robot_control_cmd", cmd_msg.encode())
            time.sleep(0.2)  # 每 0.2 秒发送一次

    except KeyboardInterrupt:
        # 捕获键盘中断，发送纯阻尼模式命令
        cmd_msg.mode = 7  # PureDamper
        cmd_msg.gait_id = 0
        cmd_msg.duration = 0
        cmd_msg.life_count += 1
        lcm_cmd.publish("robot_control_cmd", cmd_msg.encode())
        pass
    sys.exit()  # 退出程序

if __name__ == '__main__':
    main()